import geopandas as gpd
import numpy as np
import os
import gc
import math
from shapely.geometry import Polygon, MultiPolygon, Point
from shapely.validation import make_valid
from shapely.ops import unary_union
import re
from datetime import datetime

def export_forest_model_clean(geojson_path, output_obj_path, scale=100000, thickness=2.0):
    """
    优化版林场模型导出：
    1. 移除无效面、重复面和面积过小的碎面
    2. 确保每个地块只有必要的面（顶部+底部+侧面）
    3. 保持相邻地块紧密贴合，解决破碎感问题
    """
    print(f"🌲 开始导出优化面结构的林场模型...")
    print(f"📁 输入GeoJSON: {geojson_path}")
    print(f"📁 输出OBJ文件: {output_obj_path}")
    print(f"📏 厚度设置: {thickness} 单位")
    print(f"✨ 优化目标：移除无效面，解决模型破碎感")
    print(f"⏰ 开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

    # 检查输入文件
    if not os.path.exists(geojson_path):
        print(f"❌ 错误：输入文件不存在 - {geojson_path}")
        return False

    try:
        # 1. 读取并预处理数据
        gdf = gpd.read_file(geojson_path)
        print(f"✅ 成功读取GeoJSON：共包含 {len(gdf)} 个小班地块")

        # 数据清洗与标准化
        valid_plots = []
        for idx, row in gdf.iterrows():
            林场 = str(row.get("LIN_CHANGM", "未知林场")).strip()
            林班 = str(row.get("LIN_BAN", "0000")).strip().zfill(4)
            小班 = str(row.get("XIAO_BAN", "00000")).strip().zfill(5)
            几何 = row.geometry

            林场 = re.sub(r'[\\/*?:"<>|]', '', 林场)
            完整名称 = f"{林场}_{林班}_{小班}"

            # 修复无效几何
            if not 几何.is_valid:
                几何 = make_valid(几何)

            # 过滤面积过小的地块（避免碎面）
            if not 几何.is_empty and 几何.area > 1e-8:  # 面积阈值，过滤微小地块
                valid_plots.append({
                    "林场": 林场,
                    "林班": 林班,
                    "小班": 小班,
                    "完整名称": 完整名称,
                    "几何": 几何
                })
            else:
                print(f"⚠️  跳过空几何或微小地块：{完整名称}")

        if not valid_plots:
            print(f"❌ 无有效小班数据")
            return False
        print(f"✅ 有效小班数量：{len(valid_plots)}")

        # 2. 建立层级结构
        forest_hierarchy = {}
        for plot in valid_plots:
            林场 = plot["林场"]
            林班 = plot["林班"]

            if 林场 not in forest_hierarchy:
                forest_hierarchy[林场] = {}
            if 林班 not in forest_hierarchy[林场]:
                forest_hierarchy[林场][林班] = []

            forest_hierarchy[林场][林班].append(plot)

        print(f"✅ 层级结构：{len(forest_hierarchy)} 个林场 → {sum(len(lb) for lb in forest_hierarchy.values())} 个林班")

        # 3. 计算全局偏移（提高坐标精度）
        all_coords = []
        for plot in valid_plots:
            geom = plot["几何"]
            processed_geoms = process_geometry(geom)
            for g in processed_geoms:
                all_coords.extend(g.exterior.coords)
                for interior in g.interiors:
                    all_coords.extend(interior.coords)

        coords_array = np.array(all_coords, dtype=np.float64)
        min_lon, min_lat = np.min(coords_array, axis=0)
        max_lon, max_lat = np.max(coords_array, axis=0)
        global_offset = np.array([(min_lon + max_lon) / 2, (min_lat + max_lat) / 2], dtype=np.float64)
        print(f"📊 全局地理范围：\n   经度：{min_lon:.10f} ~ {max_lon:.10f}\n   纬度：{min_lat:.10f} ~ {max_lat:.10f}")

        # 4. 创建坐标映射表（确保相邻模型共享顶点）
        print(f"🔍 创建坐标映射表...")
        coord_map = create_coord_map(all_coords, precision=10)
        print(f"✅ 坐标映射表创建完成：{len(coord_map)} 个唯一坐标点")

        # 5. 模型参数配置
        COORDINATE_PRECISION = 10  # 坐标精度
        MIN_FACE_AREA = 1e-6      # 最小面面积阈值，过滤无效小面

        # 6. 生成OBJ数据
        obj_data = []
        total_vertices = 0
        total_faces = 0
        removed_faces = 0  # 统计移除的无效面数量

        # 6.1 文件头
        obj_data.append(f"# 林场模型（优化面结构版）")
        obj_data.append(f"# 生成时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        obj_data.append(f"# 特性：移除无效面和碎面 | 相邻地块紧密贴合 | 厚度 {thickness}")
        obj_data.append("")

        # 6.2 写入所有唯一顶点（顶部+底部）并保存顶点数据用于面积计算
        print(f"📝 写入顶点...")
        sorted_coords = sorted(coord_map.items(), key=lambda x: x[1])
        vertices = []  # 存储顶点坐标数据
        for (lon, lat), idx in sorted_coords:
            x, y = (lon - global_offset[0]) * scale, (lat - global_offset[1]) * scale

            # 顶部顶点 (z=0)
            top_vertex = (x, y, 0.0)
            obj_data.append(f"v {x:.{COORDINATE_PRECISION}f} {y:.{COORDINATE_PRECISION}f} 0.0")
            vertices.append(top_vertex)

            # 底部顶点 (z=-thickness)
            bottom_vertex = (x, y, -thickness)
            obj_data.append(f"v {x:.{COORDINATE_PRECISION}f} {y:.{COORDINATE_PRECISION}f} {-thickness}")
            vertices.append(bottom_vertex)

        total_vertices = len(coord_map) * 2  # 每个坐标点有2个顶点
        obj_data.append(f"\n# 顶点总数：{total_vertices}")
        obj_data.append("")

        # 6.3 生成优化后的面（移除无效面）
        print(f"🔗 生成优化后的面...")
        for 林场名, 林班字典 in forest_hierarchy.items():
            obj_data.append(f"o {林场名}")
            obj_data.append(f"usemtl Material_{林场名}")
            obj_data.append("")

            for 林班名, 小班列表 in 林班字典.items():
                林班组名 = f"{林场名}_{林班名}"
                obj_data.append(f"g {林班组名}")
                obj_data.append(f"usemtl Material_{林班组名}")
                obj_data.append("")

                for 小班 in 小班列表:
                    小班全名 = 小班["完整名称"]
                    obj_data.append(f"g {小班全名}")
                    obj_data.append(f"usemtl Material_{小班全名}")
                    obj_data.append(f"# 小班：{小班全名}")

                    geom = 小班["几何"]
                    processed_geoms = process_geometry(geom)
                    for geom_idx, g in enumerate(processed_geoms):
                        # 生成面并过滤无效面，传入vertices参数
                        faces, removed = generate_optimized_faces(
                            g, coord_map, global_offset, scale, thickness,
                            len(coord_map), min_area=MIN_FACE_AREA,
                            vertices=vertices  # 传递顶点数据
                        )
                        removed_faces += removed

                        if not faces:
                            continue

                        # 写入面
                        for face in faces:
                            obj_data.append(f"f {face[0]} {face[1]} {face[2]}")
                        total_faces += len(faces)

                    obj_data.append("")
                obj_data.append("")
            obj_data.append("")

        # 7. 写入文件
        os.makedirs(os.path.dirname(output_obj_path), exist_ok=True)
        with open(output_obj_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(obj_data))

        # 8. 输出统计
        print(f"\n🎉 优化面结构的模型导出完成！")
        print(f"📊 模型统计：")
        print(f"   - 层级：{len(forest_hierarchy)} 林场 → {sum(len(lb) for lb in forest_hierarchy.values())} 个林班 → {len(valid_plots)} 小班")
        print(f"   - 几何：{total_vertices:,} 个顶点 | {total_faces:,} 个有效面")
        print(f"   - 优化：共移除 {removed_faces} 个无效面/碎面")
        print(f"   - 厚度：{thickness} 单位 | 坐标精度：{COORDINATE_PRECISION} 位小数")
        print(f"   - 文件：{output_obj_path}")
        print(f"⏰ 完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

        gc.collect()
        return True

    except Exception as e:
        print(f"❌ 导出错误: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

def create_coord_map(coords, precision=10):
    """创建坐标映射表，确保相邻地块共享相同顶点"""
    coord_map = {}
    idx = 0
    for coord in coords:
        lon = round(coord[0], precision)
        lat = round(coord[1], precision)

        if (lon, lat) not in coord_map:
            coord_map[(lon, lat)] = idx
            idx += 1
    return coord_map

def process_geometry(geom):
    """处理几何对象，确保边界闭合和有效性"""
    processed = []
    if isinstance(geom, Polygon):
        if not geom.is_empty:
            # 确保多边形闭合
            if not np.allclose(geom.exterior.coords[0], geom.exterior.coords[-1], atol=1e-10):
                exterior_coords = list(geom.exterior.coords) + [geom.exterior.coords[0]]
                closed_geom = Polygon(exterior_coords, interiors=[list(interior.coords) for interior in geom.interiors])
                processed.append(closed_geom)
            else:
                processed.append(geom)
    elif isinstance(geom, MultiPolygon):
        for poly in geom.geoms:
            if not poly.is_empty:
                if not np.allclose(poly.exterior.coords[0], poly.exterior.coords[-1], atol=1e-10):
                    exterior_coords = list(poly.exterior.coords) + [poly.exterior.coords[0]]
                    closed_poly = Polygon(exterior_coords)
                    processed.append(closed_poly)
                else:
                    processed.append(poly)
    return processed

def generate_optimized_faces(geom, coord_map, global_offset, scale, thickness, coord_count, min_area=1e-6, precision=10, vertices=None):
    """
    生成优化后的面，移除无效面：
    1. 过滤面积过小的碎面
    2. 移除重复面和退化面（顶点共线的面）
    3. 确保每个面都是有效的三角形
    """
    # 确保顶点数据有效
    if vertices is None:
        vertices = []

    if not isinstance(geom, Polygon) or geom.is_empty:
        return [], 0

    faces = []
    removed = 0  # 记录移除的面数量
    face_set = set()  # 用于检测重复面

    # 处理外部轮廓
    exterior_coords = np.array(geom.exterior.coords)
    top_indices = []   # 顶部顶点索引（z=0）
    bottom_indices = []  # 底部顶点索引（z=-thickness）

    # 收集顶点索引
    for coord in exterior_coords:
        lon = round(coord[0], precision)
        lat = round(coord[1], precision)

        if (lon, lat) not in coord_map:
            continue

        coord_idx = coord_map[(lon, lat)]
        top_idx = coord_idx * 2 + 1    # OBJ索引从1开始
        bottom_idx = coord_idx * 2 + 2  # 底部顶点

        top_indices.append(top_idx)
        bottom_indices.append(bottom_idx)

    # 生成外部轮廓的面（并过滤无效面）
    if len(top_indices) >= 3:
        # 1. 顶部面（地表）
        for i in range(1, len(top_indices) - 1):
            # 验证索引有效性
            if top_indices[0] > len(vertices) or top_indices[i] > len(vertices) or top_indices[i+1] > len(vertices):
                removed += 1
                continue

            face = tuple(sorted([top_indices[0], top_indices[i], top_indices[i+1]]))
            # 检查是否重复
            if face in face_set:
                removed += 1
                continue
            # 计算面面积，过滤过小的面
            if calculate_face_area(face, vertices, scale) < min_area:
                removed += 1
                continue
            faces.append(face)
            face_set.add(face)

        # 2. 底部面（垂坠底面）
        for i in range(len(bottom_indices) - 2, 0, -1):
            # 验证索引有效性
            if bottom_indices[0] > len(vertices) or bottom_indices[i] > len(vertices) or bottom_indices[i+1] > len(vertices):
                removed += 1
                continue

            face = tuple(sorted([bottom_indices[0], bottom_indices[i], bottom_indices[i+1]]))
            if face in face_set:
                removed += 1
                continue
            if calculate_face_area(face, vertices, scale) < min_area:
                removed += 1
                continue
            faces.append(face)
            face_set.add(face)

        # 3. 侧面（连接顶部和底部）
        for i in range(len(top_indices) - 1):
            v1 = top_indices[i]
            v2 = top_indices[i+1]
            v3 = bottom_indices[i+1]
            v4 = bottom_indices[i]

            # 验证索引有效性
            if v1 > len(vertices) or v2 > len(vertices) or v3 > len(vertices):
                removed += 1
                continue

            # 第一个三角面
            face1 = tuple(sorted([v1, v2, v3]))
            if face1 not in face_set and calculate_face_area(face1, vertices, scale) >= min_area:
                faces.append(face1)
                face_set.add(face1)
            else:
                removed += 1

            # 验证索引有效性
            if v1 > len(vertices) or v3 > len(vertices) or v4 > len(vertices):
                removed += 1
                continue

            # 第二个三角面
            face2 = tuple(sorted([v1, v3, v4]))
            if face2 not in face_set and calculate_face_area(face2, vertices, scale) >= min_area:
                faces.append(face2)
                face_set.add(face2)
            else:
                removed += 1

    # 处理内部孔洞
    for interior in geom.interiors:
        interior_coords = np.array(interior.coords)
        hole_top = []
        hole_bottom = []

        for coord in interior_coords:
            lon = round(coord[0], precision)
            lat = round(coord[1], precision)

            if (lon, lat) not in coord_map:
                continue

            coord_idx = coord_map[(lon, lat)]
            hole_top.append(coord_idx * 2 + 1)
            hole_bottom.append(coord_idx * 2 + 2)

        # 生成孔洞的面
        if len(hole_top) >= 3:
            # 1. 顶部孔洞面
            for i in range(len(hole_top) - 2, 0, -1):
                # 验证索引有效性
                if hole_top[0] > len(vertices) or hole_top[i] > len(vertices) or hole_top[i+1] > len(vertices):
                    removed += 1
                    continue

                face = tuple(sorted([hole_top[0], hole_top[i], hole_top[i+1]]))
                if face in face_set:
                    removed += 1
                    continue
                if calculate_face_area(face, vertices, scale) < min_area:
                    removed += 1
                    continue
                faces.append(face)
                face_set.add(face)

            # 2. 底部孔洞面
            for i in range(1, len(hole_bottom) - 1):
                # 验证索引有效性
                if hole_bottom[0] > len(vertices) or hole_bottom[i] > len(vertices) or hole_bottom[i+1] > len(vertices):
                    removed += 1
                    continue

                face = tuple(sorted([hole_bottom[0], hole_bottom[i], hole_bottom[i+1]]))
                if face in face_set:
                    removed += 1
                    continue
                if calculate_face_area(face, vertices, scale) < min_area:
                    removed += 1
                    continue
                faces.append(face)
                face_set.add(face)

            # 3. 孔洞侧面
            for i in range(len(hole_top) - 1):
                v1 = hole_top[i]
                v2 = hole_top[i+1]
                v3 = hole_bottom[i+1]
                v4 = hole_bottom[i]

                # 验证索引有效性
                if v2 > len(vertices) or v1 > len(vertices) or v3 > len(vertices):
                    removed += 1
                    continue

                # 第一个三角面
                face1 = tuple(sorted([v2, v1, v3]))
                if face1 not in face_set and calculate_face_area(face1, vertices, scale) >= min_area:
                    faces.append(face1)
                    face_set.add(face1)
                else:
                    removed += 1

                # 验证索引有效性
                if v3 > len(vertices) or v1 > len(vertices) or v4 > len(vertices):
                    removed += 1
                    continue

                # 第二个三角面
                face2 = tuple(sorted([v3, v1, v4]))
                if face2 not in face_set and calculate_face_area(face2, vertices, scale) >= min_area:
                    faces.append(face2)
                    face_set.add(face2)
                else:
                    removed += 1

    # 将元组转换为列表，保持OBJ格式要求
    return [[f1, f2, f3] for (f1, f2, f3) in faces], removed

def calculate_face_area(face, vertices, scale):
    """计算三角面的面积，用于过滤过小的面"""
    try:
        # 从顶点索引获取坐标（索引减1是因为OBJ索引从1开始）
        v1 = vertices[face[0]-1]
        v2 = vertices[face[1]-1]
        v3 = vertices[face[2]-1]

        # 提取X和Y坐标
        x1, y1, _ = v1
        x2, y2, _ = v2
        x3, y3, _ = v3

        # 计算三角形面积（使用 shoelace 公式）
        area = 0.5 * abs((x2 - x1)*(y3 - y1) - (y2 - y1)*(x3 - x1))
        return area
    except:
        # 发生错误时返回0，确保会被过滤掉
        return 0

# 执行入口
if __name__ == "__main__":
    # 配置路径
    input_geojson = r"C:\Users\ma\Desktop\鹤北geojson\红松果林.geojson"
    output_obj = r"C:\Users\ma\Desktop\鹤北geojson\按林场分割模型.obj"

    # 执行导出
    success = export_forest_model_clean(
        geojson_path=input_geojson,
        output_obj_path=output_obj,
        scale=100000,
        thickness=2.0  # 可根据需要调整厚度
    )

    if success:
        print("\n🎉 优化面结构的模型导出完成！已移除无效面，模型不再有破碎感")
    else:
        print("\n❌ 导出失败")